#!/usr/bin/env python

"""
SGF -- parse an SGF file into a Game (of Nodes).

spec: http://www.red-bean.com/sgf/

Ken Harris <kengruven@gmail.com>
2006.8.30
"""

# FIXME: using a global here ('size') is bad ... but with python's
# scoping rules i'm not sure how else to do it


import game
from point import Point

## import psyco
## psyco.full()


##
## LOADING
##

from pyparsing import *
# ParserElement.enablePackrat() -- makes things much worse

# the SGF grammar:
text = QuotedString(quoteChar="[", escChar="\\", escQuote=None, multiline=True, unquoteResults=False, endQuoteChar="]")
prop = Word(srange("[A-Z]"), min=1, max=2) + OneOrMore(text)
node = Suppress(";") + OneOrMore(prop)
branch = Forward()
branch << Suppress("(") + OneOrMore(node | branch) + Suppress(")")
sgf = OneOrMore(branch)

# property
def props2dict(string, location, tokens):
    __pychecker__ = 'unusednames=string,location'
    return {tokens[0]: tokens[1:]}
prop.setParseAction(props2dict)


size = 19 # default size is always 19

# unescape \-escaped SGF string
import re
bs = re.compile("\\\\(.)")
def unescape(s):
    # remove \-escaped newlines
    s = ''.join(s.split("\\\n\r"))
    s = ''.join(s.split("\\\r\n"))
    s = ''.join(s.split("\\\n"))
    s = ''.join(s.split("\\\r")) # ... all 4 kinds!

    s = bs.sub("\\1", s)

    # note: if you use unquoteResults=True, everything seems peachy until you
    # try to unescape \\\r\n (or \\\r) -- WHY?  (it's late ... i'm not sure i understand...)

    return s

# node
def tokens2node(string, location, tokens):
    __pychecker__ = 'unusednames=string,location'
    output = {}
    for token in tokens:

        # remove []'s, since unquoteResults=False now
        for key in token:
            token[key] = [x[1:-1] for x in token[key]]

        output.update(token)

    global size

    # need size everywhere -- including parsing the root node!  so do it first, if present.
    if 'SZ' in output:
        size = int(output['SZ'][0])

    # make values into appropriate values (list of Point, etc.), then return
    for key in output:
        if key in ['B', 'W']: # points
            if output[key][0]=='' or (output[key][0]=='tt' and size<=19):
                output[key] = None
            else:
                output[key] = Point(sgf=output[key][0], size=size) # FF[4] says this can only have one point
        elif key in ['TW', 'TB', 'TR', 'SQ', 'CR', 'MA', 'AB', 'AW', 'M', 'L', 'SL', 'AE']: # point lists
            points = []
            for text in output[key]:
                if ':' in text:
                    # decode a compressed point list
                    corners = [Point(sgf=subtext, size=size) for subtext in text.split(':')]
                    assert len(corners)==2
                    row1,rown = min(corners[0].row, corners[1].row), \
                                max(corners[0].row, corners[1].row)
                    col1,coln = min(corners[0].col, corners[1].col), \
                                max(corners[0].col, corners[1].col)

                    for row in range(row1, rown+1):
                        for col in range(col1, coln+1):
                            point = Point(row=row, col=col, size=size)
                            points += [ point ]
                else:
                    point = Point(sgf=text, size=size)
                    points += [ point ]
            output[key] = points
        elif key in ['LB']: # {point:str} dict
            # input is a list of strings like ['nd:foo', 'nb:bar']
            output[key] = dict((Point(sgf=unescape(text[:2]), size=size), text[3:])
                               for text in output[key])
        elif key in ['LN', 'DD', 'AR']: # from:to pairs (not quite a compressed point list)
            # make into 2-tuple of Points
            output[key] = tuple(Point(sgf=text, size=size) for text in output[key][0].split(':'))
        else: # just a string
            assert len(output[key]) == 1 # (is this a safe assertion to make?  causes C[][] to fail...)
            output[key] = unescape(output[key][0])

        # DONE:
        # - B, W are points
        # - TW, TB, TR, SQ, CR, MA, AB are point lists
        # - LB is a {point:str} dict
        # GENERIC:
        # - it would be great if i had a dict of these, like {'LB': {Point:str}, 'B': [Point], ...}

    node = game.Node()
    node.properties = output
    return node

node.setParseAction(tokens2node)

# branch
def tokens2branch(string, location, tokens):
    __pychecker__ = 'unusednames=string,location'

    # BACKWARDS COMPATIBILITY:
    # - get a list of Node
    # - want each one's .children to point to the next one
    # oh, and branching is kind of weird in SGF:
    node = tokens[0]
    i = 1
    while i < len(tokens):
        if tokens[i].children:
            node.children += [ tokens[i] ]

            # fix up children
            for child in node.children:
                child.parent = node
                child.number = node.number + 1

                # BUG HERE: child.number isn't assigned correctly for everything.
                # perhaps it's easiest to just fix up everything at the end!

            i += 1
        else:
            node.children = [ tokens[i] ]

            # fix up children
            for child in node.children:
                child.parent = node
                child.number = node.number + 1

            node = tokens[i]
            i += 1

    return tokens[0]

branch.setParseAction(tokens2branch)


def load(filename):
    # load file/URL
    try:
        import gnomevfs
        text = gnomevfs.read_entire_file(filename) # ?? (can't do local relative filenames)
    except:
        text = file(filename).read() # but this, of course, can

    # convert from ISO-8859-1 -- this is the default encoding, by SGF FF[4]
    text = unicode(text, 'Latin-1')

    # parse
    main_branches = sgf.parseString(text)

    # make into game object
    return game.Game(main_branches[0])


##
## SAVING
##

def _escape(s):
    """escape the string |s| for an SGF [] string."""
    return s.replace("\\", "\\\\").replace("]", "\\]")

def _save_node(node, out):
    """save(print) the node |node| to file |out|."""
    out.write(";")
    for key in node.properties:
        out.write(key)
        value_or_values = node.properties[key]
        if isinstance(value_or_values, list):
            for value in value_or_values:
                try:
                    out.write("[%s]" % value.as_sgf())
                except AttributeError:
                    out.write("[%s]" % _escape(value))
        else:
            if isinstance(value_or_values, Point):
                out.write("[%s]" % value_or_values.as_sgf())
            elif value_or_values == None:
                out.write("[]")
            else:
                out.write("[%s]" % _escape(value_or_values))

def _save_branch(node, out):
    out.write("(")
    while True:
        _save_node(node, out)
        if len(node.children)==1:
            node = node.children[0]
        elif node.children:
            for child in node.children:
                _save_branch(child, out)
            break # -- awkward!
        else:
            break
    out.write(")")

def save(game, filename):
    out = open(filename, 'w')
    try:
        _save_branch(game.head, out)
    finally:
        out.close()

# ----------------------------------------
import unittest
class SGFTest(unittest.TestCase):
    def testUnescape(self):
        # backslash-escape a normal char = do nothing
        self.assertTrue(unescape("\\foo") == "foo")
        self.assertTrue(unescape("\\\\") == "\\")

        # backslash-escape a newline = remove it
        self.assertTrue(unescape("abc\\\ndef") == "abcdef")
        self.assertTrue(unescape("abc\\\rdef") == "abcdef")
        self.assertTrue(unescape("abc\\\r\ndef") == "abcdef")
        self.assertTrue(unescape("abc\\\n\rdef") == "abcdef")

if __name__=='__main__':
    unittest.main()
